package com.pig4cloud.pigx.knowledge.support.function;

import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.spring.SpringUtil;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.pig4cloud.pigx.common.core.util.R;
import com.pig4cloud.pigx.common.security.service.PigxUser;
import com.pig4cloud.pigx.knowledge.dto.BaseAiRequest;
import com.pig4cloud.pigx.knowledge.dto.ChatMessageDTO;
import com.pig4cloud.pigx.knowledge.support.annotation.FieldPrompt;
import com.pig4cloud.pigx.knowledge.support.feign.RemoteTableInfoService;
import com.pig4cloud.pigx.knowledge.support.util.PromptBuilder;
import lombok.Data;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.select.Select;
import org.springframework.stereotype.Component;

import java.util.Map;

/**
 * chat2 SQL函数调用
 *
 * @author lengleng
 * @date 2024/07/30
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class Chat2SqlFunctionCalling implements FunctionCalling<SqlRequest> {

	/**
	 * 展示
	 * @return boolean
	 */
	@Override
	public boolean showFunction() {
		return false;
	}

	/**
	 * 检查参数是否正确
	 * @param sqlRequest 请求信息
	 * @param userDetails 用户信息
	 * @param extDetails
	 * @return
	 */
	@Override
	public R<String> checkParams(SqlRequest sqlRequest, PigxUser userDetails, ChatMessageDTO.ExtDetails extDetails) {

		log.info("Chat2SQL 结果：{}", sqlRequest.getSql());

		Statement statement = null;
		try {
			statement = CCJSqlParserUtil.parse(sqlRequest.getSql());
		}
		catch (JSQLParserException e) {
			return R.failed(StrUtil.format("SQL: {} 解析失败", sqlRequest.getSql()));
		}
		if (statement instanceof Select) {
			sqlRequest.setSql(statement.toString());
			return R.ok();
		}

		return R.failed(StrUtil.format("SQL: {} 解析失败，安全问题不支持非查询语句", sqlRequest.getSql()));
	}

	/**
	 * 处理
	 * @param sqlRequest SQL 请求
	 * @param userDetails 用户详细信息
	 * @param extDetails 扩展详细信息
	 * @return {@link R }
	 */
	@Override
	public R<String> handle(SqlRequest sqlRequest, PigxUser userDetails, ChatMessageDTO.ExtDetails extDetails) {
		RemoteTableInfoService remoteTableInfoService = SpringUtil.getBean(RemoteTableInfoService.class);

		R execSQLResult = remoteTableInfoService.execSQL(extDetails.getTable().getDsName(), sqlRequest.getSql());
		// 增强一下返回数据
		if (execSQLResult.isOk()) {
			String result = PromptBuilder.render("chat2db-result.st",
					Map.of("sql", sqlRequest.getSql(), "result", execSQLResult.getData().toString()));
			return R.ok(result);
		}
		return R.failed(StrUtil.format("SQL 执行失败: {}", execSQLResult.getMsg()));
	}

}

@Data
@JsonInclude(JsonInclude.Include.NON_NULL)
class SqlRequest extends BaseAiRequest {

	@FieldPrompt("SQL statements that the database can execute")
	private String sql;

}
